import os,glob,h5py,sys,random,shutil
import SimpleITK as sitk
import numpy as np
import skimage.morphology as skimor
from scipy.spatial.distance import cdist
from scipy.ndimage.filters import gaussian_filter
import matplotlib.pyplot as plt
from PIL import Image
# ! <<< Outdated
# from keras.utils import to_categorical
import tensorflow as tf
from tensorflow.keras.utils import to_categorical
# ! >>>
sys.dont_write_bytecode = True

def dim_2_categorical(label,numclass):
	dims = label.ndim
	if dims==2:
		col, row = label.shape
		exlabel = np.zeros((numclass, col, row))
		for i in range(0,numclass):
			exlabel[i,] = np.asarray(label == i).astype(np.uint8)
	elif dims==3:
		leng,col,row = label.shape
		exlabel  = np.zeros((numclass,leng,col,row))
		for i in range(0,numclass):
			exlabel[i,] = np.asarray(label == i).astype(np.uint8)
	return exlabel

"""---------------------------------------------Min Max Normalization--------------------------------------"""
def min_max_normalization(img, lmin = None, rmax = None, dividend = None, quantile = None):
    newimg = img.copy()
    newimg = newimg.astype(np.float32)
    if quantile is not None:
        maxval = round(np.percentile(newimg,100-quantile))
        minval = round(np.percentile(newimg,quantile))
        newimg[newimg>=maxval] = maxval
        newimg[newimg<=minval] = minval

    if lmin is not None:
        newimg[newimg<lmin] = lmin
    if rmax is not None:
        newimg[newimg>rmax] = rmax

    minval = np.min(newimg)
    if dividend is None:
        maxval = np.max(newimg)
        newimg =(np.asarray(newimg).astype(np.float32) - minval)/(maxval-minval)
    else:
        newimg =(np.asarray(newimg).astype(np.float32) - minval)/dividend
    return newimg

"""----------------------------Sample image labels based on coordinates---------------------------"""
# ndim: dimension of the network
def sample_label_with_coor_2D(label, patchdims, patchcoors, nclass=2, model_dimension=2):
    samplenum = patchcoors.shape[0]
    if nclass==2:
        patches = np.empty(shape=(0,patchdims[0],patchdims[1],patchdims[2]), dtype=np.uint8)
    else:
        patches = np.empty(shape=(0,nclass,patchdims[1],patchdims[2]), dtype=np.uint8)

    for i in range(samplenum):
        middle = int((patchcoors[i,1]-patchcoors[i,0])/2) + patchcoors[i,0]
        if nclass==2:
            # ! <<< Add 3D
            if model_dimension == 2:
                currpatches = label[middle, patchcoors[i,2]:patchcoors[i,3], patchcoors[i,4]:patchcoors[i,5]]
                currpatches = currpatches.reshape([1,patchdims[0],patchdims[1],patchdims[2]])
            elif model_dimension == 3:
                offset = patchdims[0] // 2
                currpatches = label[(middle - offset):(middle + offset), patchcoors[i,2]:patchcoors[i,3], patchcoors[i,4]:patchcoors[i,5]]
                currpatches = currpatches.reshape([1, 1, patchdims[0],patchdims[1],patchdims[2]])
                return currpatches
                # ! >>>
        else:
            currpatches = label[middle,patchcoors[i,2]:patchcoors[i,3],patchcoors[i,4]:patchcoors[i,5]]
            currpatches = dim_2_categorical(currpatches,nclass)
            currpatches = currpatches.reshape([1,nclass,patchdims[1],patchdims[2]])
        patches = np.append(patches,currpatches,axis=0)

    patches = patches.astype(np.uint8)
    return patches
